import h5py
import torch
import torch.utils.data as data
import sys
sys.path.append("../Ours—new")
import utils.utils as utils


#一共 600 张图片   70%训练:840  30%测试 : 360
def PSNR(pred_img, gt_img):
    bs = pred_img.size(0)
    mse_err = (pred_img - gt_img).pow(2).sum(dim=1).view(bs, -1).mean(dim=1)
    psnr = 10 * (1 / mse_err).log10()
    return psnr.mean()
class Ballet(data.Dataset):
    def __init__(self,flag='train'):
         self.framnumber=100
         self.height = 768
         self.width = 1024
         self.traincames=[1,2,4,5,6]
         self.testcames=[3]
         self.cames=[1,2,3,4,5,6]
         self.h5path1='/data/kaixindata/3DVideos-distrib/MSR3DVideo-Ballet/ballet.h5'
         self.h5path='/data/kaixindata/3DVideos-distrib/MSR3DVideo-Ballet/ballet_vsrs.h5'
         self.flag=flag
         if flag=='train':
            self.cames1= self.traincames
         if flag=='test':
            self.cames1= self.testcames
         if flag=='all':
            self.cames1= self.cames
    def __len__(self):
        return len( self.cames1)*self.framnumber
    def __getitem__(self,index):
        cames=self.cames1[index//self.framnumber]
        fram=index%self.framnumber
        cams1='cam{}'.format(cames)
        with h5py.File(self.h5path,'r') as f:
            pictureL=torch.from_numpy(f[cams1]['img'][fram]/255).permute(2,0,1)  #名称为image 
        with h5py.File(self.h5path1,'r') as f:
            targets=torch.from_numpy(f[cams1]['target'][fram]/255).permute(2,0,1)   #名称为image

        return pictureL.unsqueeze(0).float(),targets.unsqueeze(0).float()

bb=Ballet('test')
utils.tensor_to_PIL(bb[0][0],'picture')
utils.tensor_to_PIL(bb[0][1],'target')
psnr=0
for i in range(0,len(bb)):
  psnr+=PSNR(bb[i][0],bb[i][1]).item()
print(psnr/(len(bb)-1))
